{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "c5cbf978-83c4-49a4-aadf-40e93ba4fab2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from torchvision import transforms\n",
    "from PIL import Image\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "ca1bd119-00eb-4b48-bc04-851921b6c5cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from torchvision import models\n",
    "\n",
    "# 辅助函数，一个卷积 + ReLU 激活的组合模块\n",
    "def convrelu(in_channels, out_channels, kernel, padding):\n",
    "    return nn.Sequential(\n",
    "        #2d卷积层\n",
    "        nn.Conv2d(in_channels, out_channels, kernel, padding=padding),\n",
    "        #原地激活，节省内存\n",
    "        nn.ReLU(inplace=True),\n",
    "    )\n",
    "\n",
    "class ResNetUNet(nn.Module):\n",
    "    def __init__(self, n_class):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.base_model = models.resnet18(pretrained=True)\n",
    "        # 加载预训练的ResNet-18,提高模型的效果，并将其子模块转换为列表base_layers\n",
    "        self.base_layers = list(self.base_model.children())\n",
    "\n",
    "        # 将ResNet18 的前3层打包成一个新的模块(layer0),通常：初始卷积层 + 批归一化(BN) + ReLU激活 + 最大池化层，提取低级特征（如边缘、颜色）\n",
    "        self.layer0 = nn.Sequential(*self.base_layers[:3])\n",
    "        #调整通道数或特征融合\n",
    "        self.layer0_1x1 = convrelu(64, 64, 1, 0)\n",
    "        # later1, 调取ResNet-18的模块4-5作为layer1。进一步提取特征，保持分辨率\n",
    "        self.layer1 = nn.Sequential(*self.base_layers[3:5])\n",
    "        self.layer1_1x1 = convrelu(64, 64, 1, 0)\n",
    "        #layer2-4,调取通道维度或增强非线性\n",
    "        self.layer2 = self.base_layers[5]\n",
    "        self.layer2_1x1 = convrelu(128, 128, 1, 0)\n",
    "        self.layer3 = self.base_layers[6]\n",
    "        self.layer3_1x1 = convrelu(256, 256, 1, 0)\n",
    "        self.layer4 = self.base_layers[7]\n",
    "        self.layer4_1x1 = convrelu(512, 512, 1, 0)\n",
    "        #将输入特征图的尺寸放大2倍\n",
    "        self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)\n",
    "        # 不断进行特征图融合，增强空间细节\n",
    "        # 最高层，上采样后的512听到特征图 + 编码器layer3的256通道跳跃特征- 总输入通道\n",
    "        self.conv_up3 = convrelu(256 + 512, 512, 3, 1)\n",
    "        #中间层，逐步减少通道数（512 - 512 - 128），同时恢复空间分辨率\n",
    "        self.conv_up2 = convrelu(128 + 512, 256, 3, 1)\n",
    "        self.conv_up1 = convrelu(64 + 256, 256, 3, 1)\n",
    "        self.conv_up0 = convrelu(64 + 256, 128, 3, 1)\n",
    "        #原始图像预处理，对输入图像进行两次3x3卷积，原始RGB通道数通常为3，提取低级特征（边缘、纹理）\n",
    "        self.conv_original_size0 = convrelu(3, 64, 3, 1)\n",
    "        self.conv_original_size1 = convrelu(64, 64, 3, 1)\n",
    "        #将预处理后的图像特征（64通道）与解码器最后一层输出（128通道）拼接 - 64+128=192通道，然后用3x3卷积 压缩到64通道\n",
    "        self.conv_original_size2 = convrelu(64 + 128, 64, 3, 1)\n",
    "        # 1x1卷积 将64通道映射到分类数n_class（如分割类别数），1x1卷积等效于全连接层，每个像素独立分类\n",
    "        self.conv_last = nn.Conv2d(64, n_class, 1)\n",
    "    \n",
    "    #定义输入数据如何通过网络的各个层和操作生成输出结果\n",
    "    def forward(self, input):\n",
    "        #输入预处理，提取输入图像的低级特征\n",
    "        x_original = self.conv_original_size0(input)\n",
    "        x_original = self.conv_original_size1(x_original)\n",
    "\n",
    "        #编码器（特征提取）\n",
    "        layer0 = self.layer0(input)\n",
    "        layer1 = self.layer1(layer0)\n",
    "        layer2 = self.layer2(layer1)\n",
    "        layer3 = self.layer3(layer2)\n",
    "        layer4 = self.layer4(layer3)\n",
    "\n",
    "        #解码器，特征融合于上采样\n",
    "        layer4 = self.layer4_1x1(layer4)\n",
    "        x = self.upsample(layer4)\n",
    "        layer3 = self.layer3_1x1(layer3)\n",
    "        #实现U-Net的跳跃连接（Skip Connection),将编码器的低级特征（高分辨率）于解码器的高级特征（低分辨率但语义丰富）融合\n",
    "        x = torch.cat([x, layer3], dim=1)\n",
    "        x = self.conv_up3(x) #逐渐上采样\n",
    "        x = self.upsample(x)\n",
    "        layer2 = self.layer2_1x1(layer2)\n",
    "        x = torch.cat([x, layer2], dim=1)\n",
    "        x = self.conv_up2(x)\n",
    "\n",
    "        x = self.upsample(x)\n",
    "        layer1 = self.layer1_1x1(layer1)\n",
    "        x = torch.cat([x, layer1], dim=1)\n",
    "        x = self.conv_up1(x)\n",
    "\n",
    "        x = self.upsample(x)\n",
    "        layer0 = self.layer0_1x1(layer0)\n",
    "        x = torch.cat([x, layer0], dim=1)\n",
    "        x = self.conv_up0(x)\n",
    "\n",
    "        x = self.upsample(x)\n",
    "        x = torch.cat([x, x_original], dim=1)\n",
    "        x = self.conv_original_size2(x)\n",
    "\n",
    "        out = self.conv_last(x)\n",
    "        \n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "3ec19ba1-f9c8-4dfe-b86a-185d8d0ace10",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 3, 256, 256])\n"
     ]
    }
   ],
   "source": [
    "transform = transforms.Compose([\n",
    "    transforms.Resize((256,256)),\n",
    "    transforms.ToTensor()])\n",
    "\n",
    "test_img_path = r\"C:\\Users\\jactv\\Documents\\xwechat_files\\wxid_kq34yw840or812_f008\\msg\\file\\2025-04\\Car\\imgs\\test\\img_1.jpg\"\n",
    "test_mask_path = r\"C:\\Users\\jactv\\Documents\\xwechat_files\\wxid_kq34yw840or812_f008\\msg\\file\\2025-04\\Car\\masks\\test\\img_1.jpg\"\n",
    "\n",
    "test_img = Image.open(test_img_path).convert('RGB')\n",
    "test_mask = Image.open(test_mask_path).convert('L')\n",
    "#图像变换及添加维度，（3，256，256）->（1，3，256，256）\n",
    "test_img_tensor = transform(test_img).unsqueeze(0)\n",
    "#打印转换后的张量形状\n",
    "print(test_img_tensor.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4fa01398-eb2c-46b5-bd72-0ab9d84ba690",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "AlexNet",
   "language": "python",
   "name": "alexnet"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.21"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
